"""
Redis记忆系统
存储Agent的对话历史，支持短期记忆
"""
import json
import logging
from typing import List, Dict, Any, Optional
from datetime import datetime
import redis
from django.conf import settings

logger = logging.getLogger(__name__)


class RedisMemory:
    """
    基于Redis的短期记忆系统
    
    功能：
    1. 存储对话历史
    2. 自动压缩（超过10轮对话）
    3. 支持多用户、多会话
    4. 自动过期（2小时）
    """

    def __init__(self, redis_config: Optional[Dict] = None):
        """
        初始化Redis连接
        
        Args:
            redis_config: Redis配置，默认使用Django配置
        """
        if redis_config is None:
            redis_config = getattr(settings, 'REDIS_CONFIG', {
                'host': 'localhost',
                'port': 6380,
                'db': 11,
            })

        try:
            self.client = redis.Redis(
                host=redis_config.get('host', 'localhost'),
                port=redis_config.get('port', 6380),
                db=redis_config.get('db', 11),
                decode_responses=True
            )
            # 测试连接
            self.client.ping()
            logger.info("Redis记忆系统初始化成功")
        except Exception as e:
            logger.error(f"Redis连接失败: {str(e)}")
            self.client = None

    def _get_key(self, user_id: int, session_id: str) -> str:
        """生成Redis Key"""
        return f"boss:agent:session:{user_id}:{session_id}"

    def save(
            self,
            user_id: int,
            session_id: str,
            user_input: str,
            agent_output: str,
            metadata: Optional[Dict] = None
    ) -> bool:
        """
        保存对话记录
        
        Args:
            user_id: 用户ID
            session_id: 会话ID
            user_input: 用户输入
            agent_output: Agent回复
            metadata: 额外元数据（如：工具调用记录）
        
        Returns:
            是否保存成功
        """
        if not self.client:
            logger.warning("Redis未连接，无法保存记忆")
            return False

        try:
            key = self._get_key(user_id, session_id)

            # 获取现有记忆
            existing = self.client.get(key)
            if existing:
                memory_data = json.loads(existing)
            else:
                memory_data = {
                    "session_id": session_id,
                    "user_id": user_id,
                    "created_at": datetime.now().isoformat(),
                    "messages": [],
                    "context": {}
                }

            # 添加新消息
            timestamp = datetime.now().isoformat()
            memory_data["messages"].extend([
                {
                    "role": "user",
                    "content": user_input,
                    "timestamp": timestamp
                },
                {
                    "role": "assistant",
                    "content": agent_output,
                    "timestamp": timestamp,
                    "metadata": metadata or {}
                }
            ])

            # 压缩历史（超过10轮对话）
            if len(memory_data["messages"]) > 20:  # 10轮对话=20条消息
                memory_data["messages"] = self._compress_messages(memory_data["messages"])

            # 保存到Redis，设置2小时过期
            self.client.setex(
                key,
                7200,  # 2小时
                json.dumps(memory_data, ensure_ascii=False)
            )

            logger.info(f"保存记忆成功: user={user_id}, session={session_id}")
            return True

        except Exception as e:
            logger.error(f"保存记忆失败: {str(e)}")
            return False

    def load(
            self,
            user_id: int,
            session_id: str
    ) -> List[Dict[str, str]]:
        """
        加载对话历史
        
        Args:
            user_id: 用户ID
            session_id: 会话ID
        
        Returns:
            对话历史列表，格式：[{"role": "user", "content": "..."}, ...]
        """
        if not self.client:
            return []

        try:
            key = self._get_key(user_id, session_id)
            data = self.client.get(key)

            if not data:
                logger.info(f"未找到历史记忆: user={user_id}, session={session_id}")
                return []

            memory_data = json.loads(data)
            messages = memory_data.get("messages", [])

            # 转换为LangChain格式
            history = []
            for msg in messages:
                history.append({
                    "role": msg["role"],
                    "content": msg["content"]
                })

            logger.info(f"加载记忆成功: user={user_id}, session={session_id}, 消息数={len(history)}")
            return history

        except Exception as e:
            logger.error(f"加载记忆失败: {str(e)}")
            return []

    def _compress_messages(self, messages: List[Dict]) -> List[Dict]:
        """
        压缩对话历史
        
        策略：保留最近6条消息（3轮对话），前面的生成摘要
        
        Args:
            messages: 原始消息列表
        
        Returns:
            压缩后的消息列表
        """
        if len(messages) <= 10:
            return messages

        # 保留最近6条
        recent_messages = messages[-6:]

        # 前面的消息生成摘要
        old_messages = messages[:-6]
        summary = f"[历史对话摘要] 共{len(old_messages) // 2}轮对话，用户主要咨询了招聘相关问题。"

        compressed = [
                         {
                             "role": "system",
                             "content": summary,
                             "timestamp": datetime.now().isoformat()
                         }
                     ] + recent_messages

        logger.info(f"压缩对话历史: {len(messages)} -> {len(compressed)}")
        return compressed

    def clear(self, user_id: int, session_id: str) -> bool:
        """
        清除某个会话的记忆
        
        Args:
            user_id: 用户ID
            session_id: 会话ID
        
        Returns:
            是否清除成功
        """
        if not self.client:
            return False

        try:
            key = self._get_key(user_id, session_id)
            self.client.delete(key)
            logger.info(f"清除记忆成功: user={user_id}, session={session_id}")
            return True
        except Exception as e:
            logger.error(f"清除记忆失败: {str(e)}")
            return False

    def get_context(self, user_id: int, session_id: str) -> Dict[str, Any]:
        """
        获取会话上下文
        
        Returns:
            上下文字典，如：{"selected_job_id": 123, "current_task": "screening"}
        """
        if not self.client:
            return {}

        try:
            key = self._get_key(user_id, session_id)
            data = self.client.get(key)

            if not data:
                return {}

            memory_data = json.loads(data)
            return memory_data.get("context", {})

        except Exception as e:
            logger.error(f"获取上下文失败: {str(e)}")
            return {}

    def update_context(
            self,
            user_id: int,
            session_id: str,
            context_update: Dict[str, Any]
    ) -> bool:
        """
        更新会话上下文
        
        Args:
            user_id: 用户ID
            session_id: 会话ID
            context_update: 要更新的上下文数据
        
        Returns:
            是否更新成功
        """
        if not self.client:
            return False

        try:
            key = self._get_key(user_id, session_id)
            data = self.client.get(key)

            if not data:
                return False

            memory_data = json.loads(data)
            memory_data["context"].update(context_update)

            self.client.setex(
                key,
                7200,
                json.dumps(memory_data, ensure_ascii=False)
            )

            logger.info(f"更新上下文成功: user={user_id}, session={session_id}")
            return True

        except Exception as e:
            logger.error(f"更新上下文失败: {str(e)}")
            return False
